# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

from collections import namedtuple
from coin.strategy.cheese.sizers.l2 import SpotL2Sizer, FuturesL2Sizer

L1Amount = namedtuple('L1Amount', [
    'balance_amt',
    'position_amt',
    'reserve_amt',
])


class L1Sizer(object):
  def get_l2(self, key):
    return self._l2_sizers[key]

  def get_balances(self):
    # TODO(inkyu): Separate this into balance store
    return {key: l2_sizer.get_balance() for key, l2_sizer in self._l2_sizers.items()}

  def position_amts(self, prices, balances):
    balance_amt_total = 0.
    position_amt_total = 0.
    reserve_amt_total = 0.
    for key, l2_sizer in self._l2_sizers.items():
      position_amts = l2_sizer.position_amts(prices[key], balances[key])
      if not position_amts:
        return None
      balance_amt_total += position_amts.balance_amt
      position_amt_total += position_amts.position_amt
      reserve_amt_total += position_amts.reserve_amt
    return L1Amount(
        balance_amt=balance_amt_total,
        position_amt=position_amt_total,
        reserve_amt=reserve_amt_total,
    )


class SpotL1Sizer(L1Sizer):
  def __init__(self, *, l1_config):
    self._l2_sizers = {}
    for key, l2_config in l1_config.items():
      self._l2_sizers[key] = SpotL2Sizer(l2_config)


class FuturesL1Sizer(L1Sizer):
  def __init__(self, *, l1_config):
    self._l2_sizers = {}
    for key, l2_config in l1_config.items():
      self._l2_sizers[key] = FuturesL2Sizer(l2_config)
